#ifndef MESHFOREACHCELL_HPP
#define MESHFOREACHCELL_HPP

#include "macros.hpp"
#include "typeHelpers.hpp"
#include "meshEnums.hpp"

namespace hms
{

template<WhichCells cells, bool indexOnly, bool inParallelRegion=false,
	typename MeshType, typename Callable>
void forEachCell( const MeshBase<MeshType>& mesh, Callable&& func ){

	using CellType = typename MeshType::CellType;

	if constexpr( indexOnly )
		static_assert(
			std::is_invocable_v<Callable, Index>,
			"Callable must take one parameter of type \"Index\"!"
		);
	else
		static_assert(
			std::is_invocable_v<Callable, const CellType&>,
			"Callable must take one parameter of the MeshType\'s CellType!"
		);

	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
		Index
			nx { mesh.derived().nCellsX() },
			ny { mesh.derived().nCellsY() };

		static constexpr Index ng {1};

		auto call = [&](Index i, Index j){
			if constexpr (indexOnly)
				func( (j+ng)*(nx + 2*ng) + (i+ng) );
			else
				func( mesh.derived().cell(i, j) );
		};
		if constexpr ( cells == WhichCells::all){
			OMP_PARALLEL_IF_NOT( inParallelRegion || omp_in_parallel(),
			for collapse(2) schedule(static),
			for ( Index j = -ng; j < ny+ng; ++j ){
				for ( Index i = -ng; i < nx+ng; ++i ){
					/* loop over all cells but skip corners */
					if ( ( j<0 || j>=ny ) && ( i<0 || i>=nx ) )
						continue;
					call(i, j);
				}
			})
		}
		if constexpr ( cells == WhichCells::domain ){
			OMP_PARALLEL_IF_NOT( inParallelRegion || omp_in_parallel(),
			for collapse(2) schedule(static),
			for ( Index j = 0; j < ny; ++j )
				for ( Index i = 0; i < nx; ++i )
					call(i, j);
			)
		}
		if constexpr ( cells == WhichCells::ghost ){
			OMP_PARALLEL_IF_NOT( inParallelRegion || omp_in_parallel(),
			sections,
			{
				/* first loop is across the bottom ghost cells */
				OMP_PRAGMA(section)
				for ( Index j = -ng; j < 0; ++j )
					for ( Index i = 0; i < nx; ++i )
						call(i, j);
				/* second loop is across the left and right ghost cells */
				OMP_PRAGMA( section )
				for ( Index j = 0; j < ny; ++j ){
					for ( Index i = -ng; i < 0; ++i )
						call(i, j);
					for ( Index i = nx; i < nx + ng; ++i )
						call(i, j);
				}
				/* third loop is across the top ghost cells */
				OMP_PRAGMA( section )
				for ( Index j = ny; j < ny + ng; ++j )
					for ( Index i = 0; i < nx; ++i )
						call(i, j);
			})
		}
	} else {
		auto call = [&](Index i){
			if constexpr (indexOnly)
				func(i);
			else
				func( mesh.cell(i) );
		};
		/* in case of an unstructured mesh or a StructMeshBase derivative 
		 * without simplified indexing, there are no corner cells that must be
		 * avoided */
		Index first, last;
		if constexpr ( cells == WhichCells::all ){
			first = 0;
			last = mesh.nCells();
		} else if constexpr ( cells == WhichCells::domain ){
			first = 0;
			last = mesh.nDomainCells();
		} else if constexpr ( cells == WhichCells::ghost ){
			first = mesh.nDomainCells();
			last = mesh.nCells();
		}
		OMP_PARALLEL_IF_NOT( inParallelRegion || omp_in_parallel(),
		for schedule(static),
		for ( Index i = first; i<last; ++i )
			call(i);
		)
	}
}
	
} // end namespace hms

#undef FOREACHCELL_DONT_PARALLELISE

#endif